import utils
import torch
import numpy as np
from sklearn.metrics import accuracy_score, roc_auc_score


class LinkPredictEval(object):
    def __init__(self, embed_filename, test_filename, n_node, n_embed, test_neg_filename=None, label_filename=None):
        self.embed_filename = embed_filename  # each line: node_id, embeddings(dim: n_embed)
        self.test_filename = test_filename  # each line: node_id1, node_id2
        self.test_neg_filename = test_neg_filename  # each line: node_id1, node_id2
        self.label_filename = label_filename
        self.n_node = n_node
        self.n_embed = n_embed
        self.emd = utils.read_embeddings(embed_filename, n_node=n_node, n_embed=n_embed)

    def eval_link_prediction(self):
        if self.test_neg_filename is not None:
            test_edges = utils.read_edges_from_file(self.test_filename)
            test_edges_neg = utils.read_edges_from_file(self.test_neg_filename)
            test_edges.extend(test_edges_neg)

            score_res = []
            for i in range(len(test_edges)):
                score_res.append(np.dot(self.emd[test_edges[i][0]], self.emd[test_edges[i][1]]))
            test_label = np.array(score_res)
            median = np.median(test_label)  # 按照得分中位数划分预测结果
            index_pos = test_label >= median
            index_neg = test_label < median
            test_label[index_pos] = 1
            test_label[index_neg] = 0
            true_label = np.zeros(test_label.shape)
            true_label[0: len(true_label) // 2] = 1
            accuracy = accuracy_score(true_label, test_label)
            auc_score = roc_auc_score(true_label, test_label)
        else:
            # test_edges = utils.read_edges_from_file(self.test_filename)
            test_edge_index = torch.load(self.test_filename).long()
            test_edge_index = torch.transpose(test_edge_index, 0, 1).numpy()
            test_label = torch.load(self.label_filename).long().numpy()

            score_res = []
            true_label = []

            for i in range(len(test_edge_index)):
                score_res.append(np.dot(self.emd[test_edge_index[i][0]], self.emd[test_edge_index[i][1]]))
                true_label.append(test_label[i])

            test_label = np.array(score_res)
            median = np.median(test_label)
            index_pos = test_label >= median
            index_neg = test_label < median
            test_label[index_pos] = 1
            test_label[index_neg] = 0
            accuracy = accuracy_score(true_label, test_label)
            auc_score = roc_auc_score(true_label, test_label)

        return accuracy, auc_score
